import sqlalchemy as sql
from sqlalchemy import or_, and_
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import deferred, sessionmaker, scoped_session

from sqlalchemy.orm import relation, backref


import sqlalchemy.types as types
import simplejson


Base = declarative_base()

import simplejson
class JSONType(types.TypeDecorator):
    """store a dict with keys of type ('CDS', 'mRNA', etc.) with values 
    being a list of pairs of ints ( [(start, stop), ...] locations) in a 
    single column as a string. 
    We'll never search on this, so dont need extra table.
    """
    impl = types.String

    def process_bind_param(self, value, dialect):
        return value and simplejson.dumps(value) or ''

    def process_result_value(self, value, dialect):
        return value and simplejson.loads(value) or ''

"""
class Location(Base):
    __tablename__ = "location"
    id           = sql.Column(sql.Integer, primary_key=True)
    feature_type = sql.Column(sql.String(32), unique=True, index=True)

    feature_id   = sql.Column(sql.Integer, sql.ForeignKey('feature.id'))

    feature      = relation('Feature', backref=backref('locations'))

""" 


class Feature(Base):
    __tablename__ = 'feature'

    id     = sql.Column(sql.Integer, primary_key=True)
    name   = sql.Column(sql.String(32), unique=True, index=True)

    chr    = sql.Column(sql.String(16), index=True)
    start  = sql.Column(sql.Integer, index=True)
    stop   = sql.Column(sql.Integer, index=True)

    strand = sql.Column(sql.Integer)

    locations = deferred(sql.Column(JSONType))
    info      = deferred(sql.Column(sql.String(192)))

    

    def __init__(self, start, stop, chr, name, strand, info=None):
        self.start = start
        self.stop  = stop
        assert start <= stop, "start must be < stop"

        self.strand = strand
        self.chr = chr
        self.name = name

        self.info = info or ''


    @classmethod
    def find(klass, start, stop, n=10):
        """
        """
        r = Session.query(Feature).filter(Feature.stop >= start)\
                                  .filter(Feature.start <= stop)
        return r.order_by(Feature.start).limit(n)
        
    @classmethod    
    def left(klass, position, n=10, inclusive=False):

        # if it's strict, it doesnt include overlapping features.
        # only those that are strictly left.
        col = inclusive and Feature.start or Feature.stop
        r = Session.query(Feature).filter(col < position)
        return r.order_by(col.desc()).limit(n)

    @classmethod    
    def right(klass, position, n=10, inclusive=False):

        col = inclusive and Feature.stop or Feature.start
        r = Session.query(Feature).filter(col > position)
        return r.order_by(col).limit(n)

    @classmethod
    def upstream(klass, feat, n=10, inclusive=False):
        if not hasattr(feat, 'start'):
            return Feature.left(feat, n, inclusive)
        if inclusive is False: 
            position = feat.start if feat.strand == 1 else feat.stop
        else:
            position = feat.stop if feat.strand == 1 else feat.start

        return Feature.left(position, n) if feat.strand == 1 else Feature.right(position, n)    

    @classmethod
    def downstream(klass, feat, n=None, inclusive=False):
        if not hasattr(feat, 'start'):
            return Feature.right(feat, n, inclusive)
        if inclusive is True: 
            position = feat.start if feat.strand == 1 else feat.stop
        else:
            position = feat.stop if feat.strand == 1 else feat.start

        return Feature.right(position, n) if feat.strand == 1 else Feature.right(position, n)    


    def __repr__(self):
        return "%s(%i, %i, %i)" % (self.name, self.start, self.stop, self.strand)

    def sequence(self, fasta, type=None):
        #TODO
        pass





if __name__ == "__main__":

    import os; 
    try: os.unlink('_t.db')
    except: pass

    engine = sql.create_engine('sqlite:///_t.db', echo=False)
    Session = scoped_session(sessionmaker(bind=engine))
    engine.echo = False
    Base.metadata.create_all(engine)

    f = Feature(25, 36, 'chr1', 'At2g26540', -1)
    f.locations = {'CDS': [(13282379, 13282560), (13282941, 13283127)],
               'EXON': [(13281946, 13282037),
                        (13282363, 13282560),
                        (13282941, 13283411)],
               'FIVE_PRIME_UTR': [(13281946, 13282037), (13282363, 13282378)],
               'GENE': [(13281946, 13283411)],
               'MRNA': [(13281946, 13283411)],
               'THREE_PRIME_UTR': [(13283128, 13283411)]}
    
    Session.add(f)
    Session.commit()

    f1 = Session.query(Feature).filter(Feature.name=='At2g26540').one()
    print; print f1
    print; print 'EXON' in f1.locations

    print; print f1.locations

    print; print f1.locations['CDS']
    print f1.start, f1.stop

    assert len( Feature.find(36, 36).all() )
    assert len( Feature.find(30, 30).all() )
    assert len( Feature.find(25, 25).all() )
    assert len( Feature.find(24, 24).all() ) == 0
    assert len( Feature.find(37, 37).all() ) == 0

    assert len(Feature.left(34).all()) == 0, len(Feature.left(34).all())
    assert len(Feature.left(34, inclusive=True).all()) == 1, len(Feature.left(34, inclusive=True).all())
    assert len(Feature.left(37).all()) == 1, len(Feature.left(37).all())
    
    assert len(Feature.right(34).all()) == 0, len(Feature.right(34).all())
    assert len(Feature.right(34, inclusive=True).all()) == 1, len(Feature.right(34, inclusive=True).all())

    import time
    t0 = time.time()
    for i in range(20, 20060, 1):
        f = Feature(i, i + 10, 'chr1', 'A%-3i' % i, (i % 40 and 1 or -1))
        Session.add(f)
    Session.commit()
    Feature.left(79)

    
    t = time.time()
    print
    print t - t0, "loaded", Session.query(Feature).count()
    print
    for i in range(1000):
        #Feature.left(79)
        Feature.right(900, 4)
        #Feature.right(79, 4)
    print
    print Feature.right(900, 5)


    engine.echo = True
    print time.time() - t
    f1 = Session.query(Feature).filter(Feature.name=='At2g26540').one()
    print f1.locations
